# pytorch到onnx的算子映射，用于参数和算子的转换
pytorch_to_op_type = {
    # 基础数学运算
    "add": "Add",
    "sub": "Sub",
    "mul": "Mul",
    "div": "Div",
    "pow": "Pow",
    "sqrt": "Sqrt",
    "exp": "Exp",
    "log": "Log",
    "abs": "Abs",
    "neg": "Neg",
    "clamp": "Clip",

    # 矩阵运算
    "Linear": "Gemm",
    "matmul": "MatMul",
    "transpose": "Transpose",
    "BatchNorm1d": "BatchNormalization",  # 也可以是 BatchNorm2d/BatchNorm3d
    "BatchNorm2d": "BatchNormalization",
    "BatchNorm3d": "BatchNormalization",

    # 卷积和池化
    "Conv1d": "Conv",
    "Conv2d": "Conv",
    "Conv3d": "Conv",
    "ConvTranspose1d": "ConvTranspose",
    "ConvTranspose2d": "ConvTranspose",
    "ConvTranspose3d": "ConvTranspose",
    "MaxPool1d": "MaxPool",
    "MaxPool2d": "MaxPool",
    "MaxPool3d": "MaxPool",
    "AvgPool1d": "AveragePool",
    "AvgPool2d": "AveragePool",
    "AvgPool3d": "AveragePool",
    "AdaptiveAvgPool1d": "GlobalAveragePool",
    "AdaptiveAvgPool2d": "GlobalAveragePool",
    "AdaptiveAvgPool3d": "GlobalAveragePool",

    # 激活函数
    "ReLU": "Relu",
    "Sigmoid": "Sigmoid",
    "Tanh": "Tanh",
    "LeakyReLU": "LeakyRelu",
    "Softmax": "Softmax",
    "LogSoftmax": "LogSoftmax",

    # 张量操作
    "reshape": "Reshape",
    "Flatten": "Flatten",
    "cat": "Concat",
    "split": "Split",
    "slice": "Slice",
    "gather": "Gather",
    "unsqueeze": "Unsqueeze",
    "squeeze": "Squeeze",
    "functional.pad": "Pad",

    # 损失函数
    "CrossEntropyLoss": "SoftmaxCrossEntropyLoss",
    "NLLLoss": "NegativeLogLikelihoodLoss",
    "MSELoss": "MeanSquaredError",

    # 其他常见算子
    "Dropout": "Dropout",
    "Identity": "Identity",
    "sum": "ReduceSum",
    "mean": "ReduceMean",
    "max": "ReduceMax",
    "min": "ReduceMin",
    "argmax": "ArgMax",
    "argmin": "ArgMin",
}


onnx_node_names = [
    "Gemm",                  # 对应 nn.Linear
    "Conv",                  # 对应 nn.Conv1d, nn.Conv2d, nn.Conv3d
    "MaxPool",               # 对应 nn.MaxPool1d, nn.MaxPool2d, nn.MaxPool3d
    "AveragePool",           # 对应 nn.AvgPool1d, nn.AvgPool2d, nn.AvgPool3d
    "Relu",                  # 对应 nn.ReLU
    "Sigmoid",               # 对应 nn.Sigmoid
    "Tanh",                  # 对应 nn.Tanh
    "LeakyRelu",             # 对应 nn.LeakyReLU
    "Softmax",               # 对应 nn.Softmax
    "BatchNormalization",    # 对应 nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d
    "Dropout",               # 对应 nn.Dropout
    "Gather",                # 对应 nn.Embedding
    "LSTM",                  # 对应 nn.LSTM
    "GRU",                   # 对应 nn.GRU
    "RNN",                   # 对应 nn.RNN
    "Flatten",               # 对应 nn.Flatten
    "DepthToSpace",          # 对应 nn.PixelShuffle
    "Resize",                # 对应 nn.Upsample
    "Pad",                   # 对应 nn.Pad
    "GlobalAveragePool",     # 对应 nn.AdaptiveAvgPool2d
    "GlobalMaxPool",         # 对应 nn.AdaptiveMaxPool2d
]